# 创建时间：2024/9/23上午9:14
import struct

_tag_pack = {
    int: lambda n: (0 if n >= 0 else 1, int.to_bytes(n if n >= 0 else -n, (n.bit_length() + 7) >> 3, 'big')),
    float: lambda n: (2, struct.pack('!d', n)),
    str: lambda n: (3, n.encode('utf-8')),
    bytes: lambda n: (4, n)
}
"""
打包映射, 可用范围0-7;
0: 正整数;
1: 负整数;
2: 浮点数;
3: 字符串;
4: 字节串.
"""

_tag_unpack = {
    0: lambda n: int.from_bytes(n, 'big'),
    1: lambda n: -int.from_bytes(n, 'big'),
    2: lambda n: struct.unpack('!d', n)[0],
    3: lambda n: n.decode('utf-8'),
    4: lambda n: n
}
"""
解包映射, tag可用范围0-7;
0: 正整数;
1: 负整数;
2: 浮点数;
3: 字符串;
4: 字节串.
"""


def _pack_one(data):
    # 根据数据类型选择打包函数
    func = _tag_pack.get(type(data), None)
    # 如果未返回函数则返回None
    if not func:
        raise TypeError(f"{type(data)}类型不可打包")
    # 打包数据, 返回一个tag和一个字节串
    tag, n_bytes = func(data)
    # 计算字节串的长度
    size = len(n_bytes)
    if size <= 16:
        # 如果字节串长度不大于16Byte
        # 拼接tag和字节长度放入1Byte作为tag
        result = ((tag << 4) | (size - 1)).to_bytes(1, 'big')
        result += n_bytes
    else:
        # 如果字节长度大于16Byte
        # 计算保存长度数值需要的字节数
        size_of_size = (((size - 17).bit_length() + 7) >> 3) if size - 17 else 1
        # 拼接tag和长度字节数, 保存到1Byte, 最高位置1, 作为tag
        result = ((1 << 3 | tag) << 4 | size_of_size - 1).to_bytes(1, 'big')
        # 拼接字节串长度
        result += (size - 17).to_bytes(size_of_size, 'big')
        # 拼接字节串
        result += n_bytes
    return result


def pack_data(*args):
    """
    打包数据, 把一个或一组数据转为一个字节串, 支持的数据类型有整型/浮点型/字符串/字节串.
    打包多个数据时可以以多个参数或一个列表的形式传入.
    打包时会根据每个数据的数据类型和数据长度生成一个对应的tag放在数据开头.
    算法:
    数据长度小于等于16字节时tag占1个字节, 最高位是0, 随后3位是数据类型, 低4位是数据长度-1;
    数据长度大于16字节时tag的第1个字节最高位是1, 随后3位是数据类型, 低4位是保存数据长度的字节数-1,
    该数值指定了第1个字节后几个字节内保存的是该数据的数据长度-17,
    数据长度最大占用16个字节, 因此单个数据的最大长度是2^(132)字节.
    :param args: 支持打包整型/浮点型/字符串/字节串数据, 多个数据可以传入多个参数或一个列表
    :return: 返回打包后的字节串, 如果类型错误则返回一个None
    """
    result = b''
    # 如果只有1个参数
    if len(args) == 1:
        if isinstance(args[0], list):
            args = args[0]
        else:
            try:
                result = _pack_one(args[0])
            except TypeError as e:
                print(e)
                return None
            return result
    # 多个参数或1个列表
    for data in args:
        try:
            result += _pack_one(data)
        except TypeError as e:
            print(e)
            return None
    return result


class UnknownTagError(TypeError):
    pass


def unpack_data(n_bytes):
    """
    pack_data对应的解包函数, 把一个字节串转回对应的数据, 返回一个列表.
    :param n_bytes: 解包的字节串
    :return: 返回一个列表
    """
    result: list[int | float | str | bytes] = []
    """解包结果"""
    n_size = len(n_bytes)
    """字节串长度"""
    pointer = 0
    """读取字节串的指针"""
    while pointer < n_size:
        tag_position = pointer
        # 从字节串中切出1个字节的字节tag
        b_tag, pointer = int.from_bytes(n_bytes[pointer:pointer + 1]), pointer + 1
        # 从字节tag中分离big标志、tag和数据大小
        big, tag, size = b_tag >> 7, (b_tag & 0x7f) >> 4, (b_tag & 0xf) + 1
        # 根据tag获取对应的解包函数
        func = _tag_unpack.get(tag)
        if func is None:
            raise UnknownTagError(f"tag识别错误, 错误tag的字节位置为{tag_position}")
        if big:
            # 如果使用了big标志, 在字节串中切出size指定的字节数, 作为数据长度保存到size
            size, pointer = int.from_bytes(n_bytes[pointer:pointer + size]) + 17, pointer + size
        # 根据数据长度切出数据部分
        b_data, pointer = n_bytes[pointer:pointer + size], pointer + size
        # 对数据解包, 保存到result中
        result.append(func(b_data))
    return result


def test():
    from utils import hex_dump
    n_bytes = pack_data(3513.613274687, 2, 'abcdefghijklmnopq', -(1 << 2 ** 10) + 1, b'aaaaaaaaaa')
    hex_dump(n_bytes)
    print(unpack_data(n_bytes))


if __name__ == '__main__':
    test()
